import {Annotation, START, StateGraph } from "@langchain/langgraph";
import {saveImages} from '../tool'

const State = Annotation.Root({
  questionType: Annotation<number>  // 0 定量问题 1 定性问题
})

//问题性质判断
const judge = (state:typeof State.State)=> {
    console.log("判断问题性质")
    return {
        questionType: 1
    }
}

//问题转换
const transform = (state) => {
    console.log("问题转换")
 return {}
}

const cube = (state) => {
    console.log('通过cube获取数据')
    return {}
}

const agent = (state) => {
    console.log('总结问题')
    return {

    }
}

const hasTransform = (state)=> {
    return state.questionType == 0 ? 'cube':'transform'
}

const workflow = new StateGraph(State)
.addNode("judge", judge)
.addNode("transform", transform)
.addNode('cube', cube)
.addNode('agent', agent)
.addEdge(START, "judge")
 .addConditionalEdges("judge", hasTransform, ['transform', 'cube'])
.addEdge('transform', 'cube')
.addEdge('cube', 'agent')
.addEdge("agent", "__end__")
.compile()

saveImages(workflow, 'test.png')

workflow.invoke({})